# encoding: utf-8
"""

@author: tongzhenguo

@time: 2020/6/12 上午11:29

@desc:
kafka读取 二次封装
现用于实时计算，查看日志信息；

"""
import time
import traceback

import arrow
from pykafka import KafkaClient
from pykafka.common import OffsetType


class KafkaProducer(object):
    def __init__(self, hosts, zookeeper_hosts, topics, broker_version):
        self.client = KafkaClient(hosts=hosts, broker_version=broker_version, zookeeper_hosts=zookeeper_hosts)
        self.topics = topics

    def producer_partition(self, topic):
        """
        生产者分区查看，主要查看生产消息时offset的变化
        :return:
        """
        topic = self.client.topics[topic.encode()]
        print(u"查看所有分区 {}".format(topic.patitions))
        earliest_offset = topic.earliest_available_offsets()
        print(u"获取最早可用的offset {}".format(earliest_offset))
        # 生产消息之前看看offset
        last_offset = topic.latest_available_offsets()
        print(u"最近可用offset {}".format(last_offset))
        # 同步生产消息
        p = topic.get_producer(sync=True)
        p.produce(str(time.time()).encode())
        # 查看offset的变化
        last_offset = topic.latest_available_offsets()
        print(u"最近可用offset {}".format(last_offset))

    def producer_designated_partition(self):
        """
        往指定分区写消息，如果要控制打印到某个分区，
        需要在获取生产者的时候指定选区函数，
        并且在生产消息的时候额外指定一个key
        :return:
        """
        for _topic in self.topics:
            topic = self.client.topics[_topic.encode()]
            p = topic.get_producer(sync=True)
            p.produce(str(time.time()).encode())
            print("topic %s , msg %s", _topic, str(time.time()))


class KafkaConsumer(object):
    def __init__(self, hosts, topics, group, broker_version):
        self.client = KafkaClient(hosts=hosts, broker_version=broker_version)
        self.topics = topics
        self.group = group

    def balance_consumer(self, message_handler):
        consumer_list = []
        for topic_str in self.topics:
            topic = self.client.topics[topic_str.encode()]
            consumer = topic.get_balanced_consumer(self.group, auto_offset_reset=OffsetType.EARLIEST,
                                                   auto_commit_enable=True)
            consumer_list.append(consumer)
        try:
            while True:
                for consumer in consumer_list:
                    msg = consumer.consume()
                    message_handler(msg.value)
        except Exception as e:
            now_str = arrow.now('Asia/Shanghai').format()
            print("error now : %s , msg : %s " % (now_str, repr(e)))
            traceback.print_exc()

    def simple_consumer(self, message_handler):
        consumer_list = []
        for topic_str in self.topics:
            topic = self.client.topics[topic_str.encode()]
            consumer = topic.get_simple_consumer(self.group, auto_offset_reset=OffsetType.LATEST,
                                                 auto_commit_enable=True)
            consumer_list.append(consumer)
        try:
            while True:
                for consumer in consumer_list:
                    msg = consumer.consume()
                    message_handler(msg.value)
        except Exception as e:
            now_str = arrow.now('Asia/Shanghai').format()
            print("error now : %s , msg : %s " % (now_str, repr(e)))
            traceback.print_exc()


if __name__ == "__main__":
    pass
